# segment.py
import cv2
import numpy as np
from joblib import load
import logging
import matplotlib.pyplot as plt
from sklearn.metrics import accuracy_score

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')

def apply_filters(image_path):
    """对输入图像应用多个滤波器并返回特征矩阵"""
    # 读取原始图像
    image = cv2.imread(image_path)
    h, w = image.shape[:2]
    
    # 转换为灰度图用于显示
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    # 应用滤波器
    mean_filtered = cv2.blur(image, (5, 5))
    gaussian_filtered = cv2.GaussianBlur(image, (5, 5), 0)
    
    sobelx = cv2.Sobel(image, cv2.CV_64F, 1, 0, ksize=5)
    sobely = cv2.Sobel(image, cv2.CV_64F, 0, 1, ksize=5)
    sobel_combined = cv2.magnitude(sobelx, sobely)
    
    canny_filtered = cv2.Canny(image, 100, 200)
    
    # 转换为灰度图并重塑
    img1 = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY).reshape(1, -1)
    img2 = cv2.cvtColor(mean_filtered, cv2.COLOR_BGR2GRAY).reshape(1, -1)
    img3 = cv2.cvtColor(gaussian_filtered, cv2.COLOR_BGR2GRAY).reshape(1, -1)
    img4 = cv2.cvtColor(sobel_combined.astype(np.uint8), cv2.COLOR_BGR2GRAY).reshape(1, -1)
    img5 = canny_filtered.reshape(1, -1)
    
    return np.vstack((img1, img2, img3, img4, img5)).T, h, w, gray

def get_labels_from_image(segment_image):
    """从分割图像中获取标签"""
    hsv_segment = cv2.cvtColor(segment_image, cv2.COLOR_BGR2HSV)
    
    # 定义颜色的HSV阈值
    red_lower = np.array([0, 120, 70])
    red_upper = np.array([10, 255, 255])
    green_lower = np.array([40, 40, 40])
    green_upper = np.array([80, 255, 255])
    blue_lower = np.array([110, 50, 50])
    blue_upper = np.array([130, 255, 255])
    yellow_lower = np.array([20, 100, 100])
    yellow_upper = np.array([30, 255, 255])
    
    # 创建掩膜
    red_mask = cv2.inRange(hsv_segment, red_lower, red_upper)
    green_mask = cv2.inRange(hsv_segment, green_lower, green_upper)
    blue_mask = cv2.inRange(hsv_segment, blue_lower, blue_upper)
    yellow_mask = cv2.inRange(hsv_segment, yellow_lower, yellow_upper)
    
    # 创建标签
    label = np.ones(red_mask.shape)
    label[red_mask == 255] = 1
    label[green_mask == 255] = 2
    label[blue_mask == 255] = 3
    label[yellow_mask == 255] = 4
    
    return label

def create_comparison_figure(original, true_segment, predicted_segment, accuracy):
    """创建对比图"""
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    
    # 显示原始图像
    axes[0].imshow(original, cmap='gray')
    axes[0].set_title('Original Image\nSandstone_2.png')
    axes[0].axis('off')
    
    # 显示真实分割图
    axes[1].imshow(true_segment, cmap='gray')
    axes[1].set_title('Segment\nSandstone_2_segment.png')
    axes[1].axis('off')
    
    # 显示预测分割图
    axes[2].imshow(predicted_segment, cmap='gray')
    axes[2].set_title(f'Segmentation\n(acc: {accuracy:.3f})')
    axes[2].axis('off')
    
    plt.tight_layout()
    return fig

def main():
    # 加载模型和scaler
    logging.info('加载模型...')
    model_data = load(r'D:\project8\examination_data\model.joblib')
    clf = model_data['clf']
    scaler = model_data['scaler']
    
    # 处理测试图像
    logging.info('处理测试图像...')
    image_path = r'D:\project8\examination_data\test_data\Sandstone_2.png'
    X, h, w, original_gray = apply_filters(image_path)
    
    # 特征标准化和预测
    logging.info('进行预测...')
    X_scaled = scaler.transform(X)
    prediction = clf.predict(X_scaled)
    
    # 重塑预测结果
    segmentation = prediction.reshape(h, w)
    
    # 创建输出图像
    output_image = np.zeros((h, w, 3), dtype=np.uint8)
    output_image[segmentation == 1] = [0, 0, 255]    # 红色
    output_image[segmentation == 2] = [0, 255, 0]    # 绿色
    output_image[segmentation == 3] = [255, 0, 0]    # 蓝色
    output_image[segmentation == 4] = [0, 255, 255]  # 黄色
    
    # 读取真实分割图并处理
    true_segment_img = cv2.imread(r'D:\project8\examination_data\test_data\Sandstone_2_segment.png')
    true_segment_img = cv2.resize(true_segment_img, (w, h))
    true_segment = get_labels_from_image(true_segment_img)
    
    # 计算准确率
    accuracy = accuracy_score(true_segment.reshape(-1), segmentation.reshape(-1))
    logging.info(f'分割准确率: {accuracy:.3f}')
    
    # 创建并保存对比图
    fig = create_comparison_figure(original_gray, true_segment, segmentation, accuracy)
    plt.savefig(r'D:\project8\examination_data\test_data\segmentation_comparison.png', 
                dpi=300, bbox_inches='tight')
    plt.close()
    
    # 保存预测的分割图
    cv2.imwrite(r'D:\project8\examination_data\test_data\Sandstone_2_predicted_segment.png', 
                output_image)
    
    # 显示结果
    cv2.imshow('Segmented Image', output_image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    
    logging.info('完成！')

if __name__ == '__main__':
    main()